import torch
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
from data_read import test_data
from bert_config_movie import MAX_LEN, device, model_dir
from fine_tuning_data_formatting import tokenizer
from fine_tuning import BertClass

# 定义一个自定义数据集类QueryData，继承自Dataset类
class QueryData(Dataset):
    def __init__(self, dataframe, tokenizer, max_len):
        # 初始化数据集的BertTokenizer
        self.tokenizer = tokenizer
        self.data = dataframe
        # 从DataFrame中获取查询文本列
        self.text = dataframe['query']
        #         self.targets = self.data.citation_influence_label
        # 初始化数据集的最大文本长度
        self.max_len = max_len

    # 定义数据集的长度，即数据集中样本的数量
    def __len__(self):
        return len(self.text)

    # 定义数据集的获取方法，根据索引获取具体样本
    def __getitem__(self, index):
        text = str(self.text[index])
        # 对查询文本进行分词，去除多余的空格
        text = " ".join(text.split())

        # 使用self.tokenizer.encode_plus()方法对文本进行编码，并生成模型的输入张量
        # 包含以下信息
        # input_ids：编码后的文本，表示为数字序列。
        # attention_mask：注意力掩码，用于指示哪些令牌在模型的注意力机制中被考虑，哪些令牌被掩盖。
        # token_type_ids：用于区分输入中不同句子或段落的令牌类型。在这个任务中，token_type_ids全为0，因为只有一个文本序列。
        inputs = self.tokenizer.encode_plus(
            text,  # 要编码的文本
            None,  # 在此情况下，表示不存在关联的另一个文本。在这个任务中，文本之间没有关联，因此使用None
            add_special_tokens=True,  # 表示在编码时添加特殊标记，例如[CLS]和[SEP]。
            max_length=self.max_len,  # 表示编码后的序列的最大长度。
            pad_to_max_length=True,  # 如果编码后的序列长度小于max_length，则在序列末尾填充0，使其长度与max_length相同。
            return_token_type_ids=True  # 用于区分输入中的不同句子或段落。在这个任务中，由于输入没有关联的另一个文本，token_type_ids全为0。
        )
        ids = inputs['input_ids']
        mask = inputs['attention_mask']
        token_type_ids = inputs["token_type_ids"]
        # 返回字典
        return {
            'ids': torch.tensor(ids, dtype=torch.long),
            'mask': torch.tensor(mask, dtype=torch.long),
            'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long)
            #             'targets': torch.tensor(self.targets[index], dtype=torch.float)
        }


# 测试函数，用于在测试集上进行模型预测
def test(model, testing_loader):
    res = []
    # 将BERT模型设置为评估模式
    model.eval()
    n_correct = 0
    n_wrong = 0
    total = 0
    tr_loss = 0
    nb_tr_steps = 0
    nb_tr_examples = 0
    # 使用torch.no_grad()上下文管理器，禁用梯度计算
    with torch.no_grad():
        for _, data in tqdm(enumerate(testing_loader, 0)):
            # 将输入数据的ids张量移动到指定的计算设备
            ids = data['ids'].to(device, dtype=torch.long)
            mask = data['mask'].to(device, dtype=torch.long)
            token_type_ids = data['token_type_ids'].to(device, dtype=torch.long)
            #             targets = data['targets'].to(device, dtype = torch.long)
            # 通过BERT模型进行前向传播，得到预测输出。
            outputs = model(ids, mask, token_type_ids)
            # 获取预测输出中每个样本预测得分最高的类别索引。
            big_val, big_idx = torch.max(outputs, dim=1)
            # 将预测结果添加到列表res中。
            res.extend(big_idx.tolist())

    return res


# 打印测试集的前两行数据
print(test_data.head(2))
# 创建测试数据集
data_to_test = QueryData(test_data[['query']], tokenizer, MAX_LEN)

# 测试参数设置
test_params = {'batch_size': 32,  # 批次大小
               'shuffle': False,  # 是否打乱顺序
               'num_workers': 0  # 运行的线程数量
               }
# 用于测试的数据加载器。
testing_loader_f = DataLoader(data_to_test, **test_params)
# 加载训练好的模型
model = torch.load(model_dir, map_location="cpu")
# model = BertClass()
# model.load_state_dict(torch.load(model_dir, map_location="cpu"), strict=False)

# 对测试集进行预测，得到预测结果
res = test(model, testing_loader_f)
# 根据预测结果和真实标签计算预测准确性
# 通过将模型预测的结果res与测试集中的真实标签test_data['label']逐个对应比较，并生成一个新的列表correct。
correct = [1 if pred == lab else 0 for pred, lab in zip(res, test_data['label'].tolist())]
# 统计列表correct中值为1的元素的个数（即预测正确的样本数），除以测试集的总样本数，即可得到模型在测试集上的准确率
print('accuracy on test set is - ', sum(correct) / len(test_data['label'].tolist()) * 100, '%')
